import asyncio

import bs4
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

from langchain_core.vectorstores import InMemoryVectorStore
from langchain_chroma import Chroma
# vector_store = Chroma(
#     collection_name="example_collection",
#     embedding_function=DashScopeEmbeddings(),
#     persist_directory="./chroma_db",  # Where to save data locally, remove if not necessary
# )
vector_store = Chroma(embedding_function=DashScopeEmbeddings(),persist_directory="chroma_db")

@tool(response_format="content_and_artifact")
def retrieve(query: str):
    """Retrieve information related to a query."""
    retrieved_docs = vector_store.similarity_search(query, k=1)
    serialized = "\n\n".join(
        (f"Source: {doc.metadata}\n" f"Content: {doc.page_content}")
        for doc in retrieved_docs
    )
    return serialized, retrieved_docs


from langchain_core.messages import SystemMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

llm = ChatOpenAI(
    api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="llama-4-scout-17b-16e-instruct",  # 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    # other params...
)
# Step 1: Generate an AIMessage that may include a tool-call to be sent.
def query_or_respond(state: MessagesState):
    """Generate tool call for retrieval or respond."""
    llm_with_tools = llm.bind_tools([retrieve])
    response = llm_with_tools.invoke(state["messages"])
    # MessagesState appends messages to state instead of overwriting
    return {"messages": [response]}


# Step 2: Execute the retrieval.
tools = ToolNode([retrieve])


# Step 3: Generate a response using the retrieved content.
async def generate(state: MessagesState):
    """Generate answer."""
    # Get generated ToolMessages
    recent_tool_messages = []
    for message in reversed(state["messages"]):
        if message.type == "tool":
            recent_tool_messages.append(message)
        else:
            break
    tool_messages = recent_tool_messages[::-1]

    # Format into prompt
    docs_content = "\n\n".join(doc.content for doc in tool_messages)
    system_message_content = (
        "You are an assistant for question-answering tasks. "
        "Use the following pieces of retrieved context to answer "
        "the question. If you don't know the answer, say that you "
        "don't know. Use three sentences maximum and keep the "
        "answer concise."
        "\n\n"
        f"{docs_content}"
    )
    conversation_messages = [
        message
        for message in state["messages"]
        if message.type in ("human", "system")
        or (message.type == "ai" and not message.tool_calls)
    ]
    prompt = [SystemMessage(system_message_content)] + conversation_messages

    # Run
    # response = llm.invoke(prompt)
    # return {"messages": [response]}
    # 使用 stream 模式逐步输出 token
    async for chunk in llm.astream(prompt):
        print(chunk.content, end="\\n", flush=True)



# Build graph
graph_builder = StateGraph(MessagesState)

graph_builder.add_node(query_or_respond)
graph_builder.add_node(tools)
graph_builder.add_node(generate)

graph_builder.set_entry_point("query_or_respond")
graph_builder.add_conditional_edges(
    "query_or_respond",
    tools_condition,
    {END: END, "tools": "tools"},
)
graph_builder.add_edge("tools", "generate")
graph_builder.add_edge("generate", END)

memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)
#
# graph_png = graph.get_graph().draw_mermaid_png()
# with open("graph.png", "wb") as f:
#     f.write(graph_png)
# #下载图片
# print(graph)

# Specify an ID for the thread
config = {"configurable": {"thread_id": "abc123"}}

input_message = "What is JAVA?"

import asyncio

async def main():
    input_message = "What is JAVA?"

    # 流式执行并输出
    async for event in graph.astream(
        {"messages": [{"role": "user", "content": input_message}]},
        stream_mode="values",
        config={"configurable": {"thread_id": "abc123"}},
    ):
        # 打印最新消息
        if event["messages"] and event["messages"][-1].type == "ai":
            print(event["messages"][-1].content, end="\\n", flush=True)

asyncio.run(main())
